(* This file is part of IsoType.
 *
 * IsoType is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * IsoType is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with IsoType.  If not, see <http://www.gnu.org/licenses/>.
 *)

type name = string
type tvar = int

module OrderedTVar =
struct
  type t = tvar
  let compare = Pervasives.compare
end

module TVarSet = Set.Make(OrderedTVar)

type t = Predefined of t list * Ast.qname
	 | Var of tvar
	 | Arrow of t * t
	 | Pair of t * t 
	 | Record of field list
	 | Algebraic of constructor list
	 | Rec of tvar * t
and field = Field of name * t
and constructor = Constructor of name * t

let unit_type = Predefined ([], Ast.Id "unit")

let is_unit = 
  function 
      Predefined ([], Ast.Id "unit") -> true
    | _ -> false 

let occurs v =
  let rec aux =
    function
	Predefined (typs, _) -> aux_list typs
      | Var v' -> v = v'
      | Arrow (typ1, typ2) -> aux typ1 || aux typ2
      | Pair (typ1, typ2) -> aux typ1 || aux typ2
      | Record fields -> List.exists aux_field fields
      | Algebraic ctors -> List.exists aux_ctor ctors
      | Rec (v', _) when v = v' -> false
      | Rec (_, typ) -> aux typ
  and aux_list typs = List.exists aux typs
  and aux_field (Field (_, typ)) = aux typ
  and aux_ctor (Constructor (_, typ)) = aux typ
  in
    aux

let fv typ = 
  let collect = List.fold_left TVarSet.union TVarSet.empty in
  let rec aux =
    function 
	Var tvar -> TVarSet.singleton tvar
      | Predefined (typs, _) -> collect (List.map aux typs)
      | Arrow (typ1, typ2) -> TVarSet.union (aux typ1) (aux typ2)
      | Pair (typ1, typ2) -> TVarSet.union (aux typ1) (aux typ2)
      | Record fields -> collect (List.map aux_field fields)
      | Algebraic ctors -> collect (List.map aux_ctor ctors)
      | Rec (tvar, typ) -> TVarSet.remove tvar (aux typ)
  and aux_field (Field (_, typ)) = aux typ
  and aux_ctor (Constructor (_, typ)) = aux typ
  in 
    TVarSet.elements (aux typ)

let rec subst tvar typ = 
  let rec aux = 
    function
      |	Predefined (typs, qname) -> Predefined (List.map aux typs, qname)
      | Var tvar' when tvar = tvar' -> typ
      | Var _ as typ' -> typ'
      | Arrow (typ1, typ2) -> Arrow (aux typ1, aux typ2)
      | Pair (typ1, typ2) -> Pair (aux typ1, aux typ2)
      | Record fields -> Record (List.map aux_field fields)
      | Algebraic ctors -> Algebraic (List.map aux_ctor ctors)
      | Rec (tvar', _) as typ' when tvar = tvar' -> typ'
      | Rec (tvar', typ') -> Rec (tvar', aux typ')
  and aux_field (Field (name, typ)) = Field (name, aux typ)
  and aux_ctor (Constructor (name, typ)) = Constructor (name, aux typ)
  in
    aux

let substs = List.fold_right (FunAux.uncurry subst)

let rec flatten_pairs =
  function
      Pair (typ1, typ2) -> typ2 :: flatten_pairs typ1
    | typ -> [typ]

let flatten_arrows = 
  let rec aux typs =
    function
	Arrow (typ1, typ2) -> aux (typ1 :: typs) typ2
      | typ -> (typs, typ)
  in
    aux []

let string_of_var x = "`" ^ string_of_int x

let rec string_of_type =
  function
      Predefined (typs, name) -> string_of_reference typs name
    | Var var -> string_of_var var
    | Arrow (typ1, typ2) -> "(" ^ string_of_type typ1 ^ " -> " ^ string_of_type typ2 ^ ")"
    | Pair (typ1, typ2) -> "(" ^ string_of_type typ1 ^ " * " ^ string_of_type typ2 ^ ")"
    | Record fields -> "{" ^ String.concat ", " (List.map string_of_field fields) ^ "}"
    | Algebraic ctors -> "[" ^ String.concat " | " (List.map string_of_ctor ctors) ^ "]"
    | Rec (var, typ) -> "FIX " ^ string_of_var var ^ "." ^ string_of_type typ
and string_of_field (Field (name, typ)) =
  name ^ " : " ^ string_of_type typ
and string_of_ctor (Constructor (name, typ)) =
  match typ with
      typ when is_unit typ -> name
    | _ -> name ^ " of " ^ string_of_type typ
and string_of_reference typs name =
  match typs with
      [] -> Pp.string_of_qname name
    | [typ] -> string_of_type typ ^ " " ^ Pp.string_of_qname name
    | typs -> "(" ^ String.concat ", " (List.map string_of_type typs) ^ ") " ^ Pp.string_of_qname name

let get_var =
  let next_var_id = ref 0 in
    fun () ->
      let id = !next_var_id in
	incr next_var_id;
	id

let translate flags types_tbl typ0 =
  let env = Hashtbl.create 211 in
  let var_env = Hashtbl.create 211 in

  let rec aux subst =
    function
	Ast.Var x -> aux_var subst x
      | Ast.Arrow (typ1, typ2) -> Arrow (aux subst typ1, aux subst typ2)
      | Ast.Tuple typs ->
	  begin  
	    match (aux_list subst typs) with
		[] -> assert false (* IMPOSSIBLE *)
	      | hd :: tl ->  List.fold_left (fun a b -> Pair(a, b)) hd tl
	  end 
      | Ast.Ref (typs, x) -> aux_ref subst (aux_list subst typs) x

  and aux_var subst x =
    try
      List.assoc x subst
    with
	Not_found ->
	  try
	    Hashtbl.find var_env x
	  with
	      Not_found ->
		let v = Var (get_var ()) in
		  Hashtbl.add var_env x v;
		  v

  and aux_ref subst typs x =
    let key = (typs, x) in
      try
	Hashtbl.find env key
      with
	  Not_found ->
	    begin
	      try
		match Hashtbl.find types_tbl x with
		    (vars, Ast.Abstract) ->
		      let res = Predefined (typs, x) in
			Hashtbl.add env key res;
			res
		  | (vars, body) ->
		      assert (List.length vars = List.length typs);
		      let v = get_var () in
		      let subst' = List.combine vars typs in
			Hashtbl.add env key (Var v);
			let body' = aux_body subst' body in
			let res = if occurs v body' then Rec (v, body') else body' in
			  Hashtbl.replace env key res;
			  res
	      with
		  Not_found ->
		    if Flags.debug flags then 
		      prerr_endline ("*********** Cannot find definition for type " ^ Pp.string_of_qname x ^ " (ASSUMING ABSTRACT)");
		    let res = Predefined (typs, x) in
		      Hashtbl.add env key res;
		      res
	    end
  and aux_body subst =
    function
	Ast.Abstract -> assert false (* IMPOSSIBLE *)
      | Ast.Alias typ -> aux subst typ
      | Ast.Algebraic ctors -> Algebraic (List.map (aux_ctor subst) ctors)
      | Ast.Record fields -> Record (List.map (aux_field subst) fields)
  and aux_ctor subst (Ast.Constructor (name, typs)) =
    match aux_list subst typs with 
	[] -> Constructor (name, unit_type)
      | hd :: tl -> Constructor (name, (List.fold_left (fun a b -> Pair(a, b)) hd tl))
  and aux_field subst (Ast.Field (name, typ)) =
    Field (name, aux subst typ)
  and aux_list subst = List.map (aux subst)
  in
    aux [] typ0

(* Normal Form && Split Normal Form *)

let nf flags =
  let rec aux =
    function 
	Pair (x, y) -> aux_irr (Pair (aux x, aux y))
      | Arrow (x, y) -> aux_irr (Arrow (aux x, aux y))
      | Record fields -> Record (List.map aux_field fields)
      | Algebraic ctors -> Algebraic (List.map aux_ctor ctors)
      | Rec (tvar, typ) -> Rec (tvar, aux typ)
      | Predefined (typs, name) -> Predefined (List.map aux typs, name)
      | typ -> typ
  and aux_irr = 
    function 
	(* Arrow (x, y) when is_unit y -> y      UNCOMMENT FOR FULL ISO *)
      | Arrow (x, y) when is_unit x -> y
      | Pair (x, Pair (y, z)) -> aux_irr (Pair (aux_irr (Pair (x, y)), z))
      | Pair (x, y) when is_unit x -> y
      | Pair (y, x) when is_unit x -> y
      | Arrow (Pair (x, y), z) ->
	  aux_irr (Arrow (x, (aux_irr (Arrow (y, z)))))
      | Arrow (x, Pair (y, z)) -> 
	  Pair (
	    aux_irr (Arrow (x, y)), 
	    aux_irr (Arrow (x, z)) 
	  )
      | typ -> typ
  and aux_field (Field (name, typ)) =
    Field (name, aux typ)
  and aux_ctor (Constructor(name, typ)) =
    Constructor (name, aux typ)
  in
    aux

let snf flags typ =
  let typ' = nf flags typ in
  let coords = flatten_pairs typ' in
  let coords' =
    List.map
      (fun typ ->
	 let subst = List.map (fun tvar -> (tvar, Var (get_var ()))) (fv typ) in
	   substs subst typ)
      coords
  in
  let res =
    match coords' with
	[] -> assert false (* IMPOSSIBLE *)
      | typ :: typs -> List.fold_left (fun typ1 typ2 -> Pair (typ1, typ2)) typ typs
  in
    (* prerr_endline (Printf.sprintf "TYPE %s\nSNF  %s" (string_of_type typ) (string_of_type res)); *)
    res
